import numpy
import json


class Accumulator:
    def __init__(self):
        pass

    def add(self, *args):
        pass

    def __str__(self):
        pass

    def json(self):
        pass


class SegClassIouAccumulator(Accumulator):

    def __init__(self, num_classes):  # num_classes: include background class; dt,gt : index 0 is background class
        super().__init__()
        self.matrix = numpy.zeros((num_classes, num_classes), dtype=numpy.int32)
        self.num_classes = num_classes

    def update_matrix(self, d, g):
        self.matrix += numpy.bincount(self.num_classes * g + d,
                                      minlength=self.num_classes ** 2).reshape(self.num_classes, self.num_classes)

    def add(self, dt, gt):
        assert dt.dtype == numpy.uint32 and dt.dtype == numpy.uint32
        for d, g in zip(dt, gt):
            self.update_matrix(d.flatten(), g.flatten())

    def json(self):
        acc = numpy.diag(self.matrix).sum() / self.matrix.sum()
        class_acc = numpy.diag(self.matrix) / self.matrix.sum(axis=1)

        class_iou = numpy.diag(self.matrix) / (
                self.matrix.sum(axis=0) + self.matrix.sum(axis=1) - numpy.diag(self.matrix))
        iou = numpy.nanmean(class_iou)
        return {'accuracy': acc.item(0), 'class_accuracy': class_acc.tolist(), 'class_iou': class_iou.tolist(),
                'iou': iou.item(0)}

    def __str__(self):
        return json.dumps(self.json())

    def reset(self):
        self.matrix[::] = 0

# if __name__ == '__main__':
#     accumulator = SegClassIouAccumulator(num_classes=3)
#     dt = numpy.array([[1, 2], [0, 2]], dtype=numpy.int32).reshape((1, 2, 2))
#     gt = numpy.array([[1, 2], [0, 2]], dtype=numpy.int32).reshape((1, 2, 2))
#     accumulator.add(dt, gt)
#     print(accumulator)
